import os
import pandas as pd
import numpy as np
from time import time
from tqdm import tqdm
import gc
import warnings

# 忽略警告信息
warnings.filterwarnings("ignore")

# 设置随机种子，保证结果可复现
np.random.seed(1997)

# 构建训练数据样本
def construct_samples_train(train_click, valid_num=20000, max_len=19, model_name='srgnn'):
    # 按照点击时间戳对训练数据进行排序
    train_click.sort_values("click_timestamp", inplace=True)

    # 加载验证集用户列表（从文件中加载）
    valid_users = np.load('valid_users.npy')
    # 筛选出验证集点击记录
    valid_click = train_click[train_click.user_id.isin(valid_users)]
    # 剩余数据作为训练集
    train_click = train_click[~train_click.user_id.isin(valid_users)]

    # 初始化训练集用户、序列和目标列表
    train_users, train_seqs, train_targets = [], [], []
    for user, val in tqdm(train_click.groupby('user_id')):
        hist_item = val['click_article_id'].values.tolist()  # 用户的点击文章列表
        # hist_time = val['click_timestamp'].values.tolist()  # 可选：用户点击时间戳列表（暂时未用到）
        hist_len = len(hist_item)  # 用户点击的历史记录长度
        if hist_len == 1:
            continue  # 如果只有一个点击记录，跳过
        elif hist_len <= max_len + 1:
            for i in range(hist_len - 1):
                # 按照历史记录长度生成训练样本
                train_users.append(user), train_seqs.append(hist_item[:i + 1]), train_targets.append(hist_item[i + 1])
        else:
            for i in range(hist_len - max_len):
                train_users.append(user), train_seqs.append(hist_item[i:i + max_len]), train_targets.append(hist_item[i + max_len])

    # 初始化验证集用户、序列和目标列表
    valid_users, valid_seqs, valid_targets = [], [], []
    for user, val in tqdm(valid_click.groupby('user_id')):
        hist_item = val['click_article_id'].values.tolist()  # 用户的点击文章列表
        # hist_time = val['click_timestamp'].values.tolist()  # 可选：用户点击时间戳列表（暂时未用到）
        hist_len = len(hist_item)  # 用户点击的历史记录长度
        if hist_len == 1:
            continue  # 如果只有一个点击记录，跳过
        elif hist_len <= max_len + 1:
            for i in range(hist_len - 2):
                # 按照历史记录长度生成训练样本
                train_users.append(user), train_seqs.append(hist_item[:i + 1]), train_targets.append(hist_item[i + 1])
            # 最后一项作为验证集目标
            valid_users.append(user), valid_seqs.append(hist_item[:-1]), valid_targets.append(hist_item[-1])
        else:
            for i in range(hist_len - max_len - 1):
                train_users.append(user), train_seqs.append(hist_item[i:i + max_len]), train_targets.append(hist_item[i + max_len])
            # 最后一项作为验证集目标
            valid_users.append(user), valid_seqs.append(hist_item[-(max_len + 1):-1]), valid_targets.append(hist_item[-1])

    # 设置保存路径
    save_path = '../samples_tmp/'
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # 保存训练集和验证集样本
    train = (train_users, train_seqs, train_targets)
    np.savez(save_path + 'train_{}.npz'.format(model_name), users=train_users, seqs=train_seqs, targets=train_targets)
    del train_users, train_seqs, train_targets  # 清理内存

    valid = (valid_users, valid_seqs, valid_targets)
    np.savez(save_path + 'valid_{}.npz'.format(model_name), users=valid_users, seqs=valid_seqs, targets=valid_targets)
    del valid_users, valid_seqs, valid_targets  # 清理内存

    # 执行垃圾回收，释放内存
    gc.collect()

    print(len(train[0]), len(valid[0]))

    return train, valid

# 构建测试数据样本
def construct_samples_test(train, test_click, max_len=19):
    # 按照点击时间戳对测试数据进行排序
    test_click.sort_values("click_timestamp", inplace=True)
    train_users, train_seqs, train_targets = list(train[0]), list(train[1]), list(train[2])
    test_users, test_seqs = [], []
    for user, val in tqdm(test_click.groupby('user_id')):
        hist_item = val['click_article_id'].values.tolist()  # 用户的点击文章列表
        # hist_time = val['click_timestamp'].values.tolist()  # 可选：用户点击时间戳列表（暂时未用到）
        hist_len = len(hist_item)  # 用户点击的历史记录长度
        if hist_len == 1:
            test_users.append(user), test_seqs.append(hist_item)  # 如果只有一个点击记录，直接加入测试集
        elif hist_len == 2:
            train_users.append(user), train_seqs.append(hist_item[:-1]), train_targets.append(hist_item[-1])  # 第一条作为训练样本
            test_users.append(user), test_seqs.append(hist_item)  # 第二条作为测试样本
        elif hist_len <= max_len + 1:
            for i in range(hist_len - 1):
                train_users.append(user), train_seqs.append(hist_item[:i + 1]), train_targets.append(hist_item[i + 1])
            test_users.append(user), test_seqs.append(hist_item)  # 余下部分作为测试样本
        else:
            for i in range(hist_len - max_len):
                train_users.append(user), train_seqs.append(hist_item[i:i + max_len]), train_targets.append(hist_item[i + max_len])
            test_users.append(user), test_seqs.append(hist_item[-max_len:])  # 仅保留最大长度的历史点击记录作为测试样本

    # 将训练集和测试集保存为元组
    train = (train_users, train_seqs, train_targets)
    del train_users, train_seqs, train_targets  # 清理内存

    test = (test_users, test_seqs)
    del test_users, test_seqs  # 清理内存
    gc.collect()

    print(len(train[0]), len(test[0]))

    return train, test
